from openai import OpenAI
from typing import Optional, Dict, Any, Literal


def generate_image(
    prompt: str,
    api_key: str,
    model: str = "doubao-seedream-3-0-t2i-250415",
    size: str = "720x1280",
    response_format: Literal["url", "b64_json"] = "url",
    base_url: str = "https://ark.cn-beijing.volces.com/api/v3",
) -> Dict[str, Any]:
    """
    Generate image using OpenAI-compatible API (Ark/Doubao).

    Args:
        prompt: Text description of the image to generate
        model: Model to use for generation (default: doubao-seedream-3-0-t2i-250415)
        size: Image size (default: "1024x1024")
        response_format: Response format - "url" or "b64_json" (default: "url")
        api_key: API key (if None, reads from ARK_API_KEY env var)
        base_url: API base URL (default: Ark Beijing endpoint)

    Returns:
        Dict containing the API response with image data

    Raises:
        Exception: If API call fails or API key is missing
    """
    if not api_key:
        raise ValueError("API key is required.")

    # Initialize client
    client = OpenAI(
        base_url=base_url,
        api_key=api_key,
    )

    # Generate image
    response = client.images.generate(
        model=model,
        prompt=prompt,
        size=size,
        response_format=response_format,
        extra_body={"watermark": False},
        extra_headers={"x-ark-moderation-scene": "skip-ark-moderation"},
    )

    return response.model_dump()


def generate_image_url(prompt: str, **kwargs) -> str:
    """
    Generate image and return the URL directly.

    Args:
        prompt: Text description of the image to generate
        **kwargs: Additional arguments to pass to generate_image()

    Returns:
        str: URL of the generated image
    """
    response = generate_image(prompt, **kwargs)
    return response["data"][0]["url"]


def generate_image_base64(prompt: str, **kwargs) -> str:
    """
    Generate image and return base64 encoded data directly.

    Args:
        prompt: Text description of the image to generate
        **kwargs: Additional arguments to pass to generate_image()

    Returns:
        str: Base64 encoded image data
    """
    kwargs["response_format"] = "b64_json"
    response = generate_image(prompt, **kwargs)
    return response["data"][0]["b64_json"]
